'''
本代码是将获取到的调和常数，插值到六边形网格，得到六边形栅格化调和常数值
然后，根据等值线提取方法，提取同潮时线以及同振幅线，实现同潮图的绘制
'''

import json
import os
import sys
import numpy as np

root_path = r"__file__"                                                                                     # 添加pydggs路径
if root_path not in sys.path:
    sys.path.insert(0, root_path)

from PyDGGS.PyDGGS import *
from ExtractOceanTide.ExtractOceanTide import hg_output_path   
from OceanSurface4D.common import level, grid_type, element_type,wave_names, H_min, H_step, G_min, G_step                         

input_dir = os.path.join(os.path.dirname(r"__file__"))                                                      # 数据位置
output_dir = os.path.join(input_dir, 'output','cotidechart')                                                # 输出数据位置

def getHgData(level, grid_type, element_type, hg_file, prop_names):
    """
    获取HG文件中指定属性名对应的值，并按照的值重新排列格网ID
    
    :param level:               离散格网等级
    :param grid_type:           离散格网类型
    :param element_type:        离散格网元素类型
    :param hg_file:             调和分析得到的迟角振幅文件
    :param prop_names:          要获取的分潮H和/或G，如['M2_H', 'M2_G', 'S2_H', 'S2_G']
    :return:                    每个格网单元的分潮H和/或G的值，如{'M2_H': {grid_id1: value1, ...},  {'S2_H': {grid_id1: value2, ...}, ...}
    """
    
    # 给格网赋值,处理振幅插值
    kernel = CreateKernelObj(grid_type)

    with open(hg_file, 'r', encoding='UTF-8') as hg_input:
        grid_hg = json.load(hg_input)

    grid_values = {}
    for prop in prop_names:
        grid_values[prop] = {}

    for cell in grid_hg['features']:
        current_id = cell['properties']['codeId']
        current_code = DGGSElementCode()
        current_code.level = level
        current_code.elementType = element_type
        id2Code(current_id, current_code, grid_type)

        current_coord = GeoCoord()
        kernel.decode(current_code, current_coord)

        cell_values = {}
        for prop in prop_names:
            cell_values[prop] = cell['properties'][prop]

        if not np.any(np.isnan(list(cell_values.values()))):
            for prop in prop_names:
                grid_values[prop][current_id] = cell_values[prop]

    return grid_values

class Contour(object):
    """
    计算全球离散格网下迟角和振幅的等值线
    """
    def __init__(self, level, grid_type, element_type, grid_prop_values,
                 min_contour_value, contour_step, is_cycle=False, cycle_bounds=None):
        """
        初始化
        
        :param level:                       离散格网等级
        :param grid_type:                   离散格网类型
        :param element_type:                离散格网元素类型
        :param grid_prop_values:            每个格网单元对应的迟角或振幅的值 {grid_id1: value1, grid_id2:value2, ...}
        :param min_contour_value:           等值线的起始值
        :param contour_step:                等间距
        :param is_cycle:                    如果计算迟角的等值线时，应为True，振幅时为False
        :param cycle_bounds:                如果计算迟角的等值线时，应为[0,360]，振幅时为None
        """
        self.kernel = CreateKernelObj(grid_type)
        self.grid_type = grid_type
        self.grid_level = level
        self.grid_element_type = element_type
        self.contour_step = contour_step
        self.min_contour_value = min_contour_value

        if not grid_prop_values:
            print("没有有效值")
            return None
        
        self.is_cycle = is_cycle
        if is_cycle:
            self.prop_data = {}
            for cell_id, cell_value in grid_prop_values.items():
                self.prop_data[cell_id] = self.wrapToCycle(cell_value, cycle_bounds)
        else:
            self.prop_data = grid_prop_values

        max_grid_value = max(self.prop_data.values())

        self.contour_list = [min_contour_value]
        next_contour_value = min_contour_value + contour_step
        while next_contour_value < max_grid_value:
            self.contour_list.append(next_contour_value)
            next_contour_value += contour_step

    def wrapToCycle(self, value, cycle_bounds):
        """
        将一个值value约束到指定周期范围内
        
        :param value:           值
        :param cycle_bounds:    存储周期的下界和上界的元组
        :return:                约束后的值
        """
        cycle_length = cycle_bounds[1] - cycle_bounds[0]

        return (value - cycle_bounds[0]) % cycle_length + cycle_bounds[0]

    def getValueClass(self, value):
        """
        获取值value对应的类别
        
        :param value:   值
        :return:        类别ID
        """
        value_class = 0

        if self.is_cycle:
            if self.contour_list[0] <= value < self.contour_list[-1]:
                value_class = 1
                for contour_value in self.contour_list[1:]:
                    if value < contour_value:
                        break

                    value_class += 1
        else:
            for contour_value in self.contour_list:
                if value < contour_value:
                    break

                value_class += 1

        return value_class

    def classifyGrid(self):
        """
        对格网按照等值线分类

        :return (classified_grid_by_id, classified_grid_by_class):  
            classified_grid_by_id:  按格元ID组织的分类后的格元  
                {grid_id1: class1, grid_id2: class2, ...}
            classified_grid_by_class: 按类别ID组织的分类后的格元 
                {class1: [grid_id1, grid_id2], class2: [grid_id3, grid_id4], ...}
            
        """
        classified_grid_by_id = {}          # 按格元ID组织的分类后的格元
        classified_grid_by_class = {}       # 按类别ID组织的分类后的格元F

        for cell_id, cell_value in self.prop_data.items():
            value_class = self.getValueClass(cell_value)
            classified_grid_by_id[cell_id] = value_class

            if value_class not in classified_grid_by_class:
                classified_grid_by_class[value_class] = []

            classified_grid_by_class[value_class].append(cell_id)

        return classified_grid_by_id, classified_grid_by_class

    def calGridCountour(self, classified_grid):
        """
        基于分类后的格网得到每一条等值线对应的格网单元

        :param classified_grid:     Contour.classifyGrid()的结果
        :return contour:            每一条等值线对应的格网单元Id    
                                    {class_id1: [grid_id1, grid_id2], class_id2: [grid_id3, grid_id4], ...}
        """
        classified_grid_by_id, classified_grid_by_class = classified_grid
        visited_satisfied_cells = {}
        for c_class, cell_ids in classified_grid_by_class.items():
            temp_visited_satisfied_cells = set()
            for eid in cell_ids:
                eid -= 64*64
                code = DGGSElementCode()
                code.level = self.grid_level
                code.elementType = self.grid_element_type
                id2Code(eid, code, self.grid_type)

                neighbor_cells = QueryResult()    #一层扩散
                self.kernel.query(DGGSTopoType_RelativeCells, code, neighbor_cells)

                for neighbor_cell_idx in range(neighbor_cells.getSize()):
                    neighbor_cell_code = neighbor_cells.getGridCodeByIndex(neighbor_cell_idx)
                    neighbor_cell_fid = code2Id(neighbor_cell_code, self.grid_type)

                    if neighbor_cell_fid not in classified_grid_by_id:
                        continue

                    neighbor_cell_class = classified_grid_by_id[neighbor_cell_fid]
                    if neighbor_cell_class != c_class:
                        temp_visited_satisfied_cells.add(eid)
                        break

            visited_satisfied_cells[c_class] = list(temp_visited_satisfied_cells)     #将同一个class下的边界点保存在一起

        return visited_satisfied_cells

def saveCotideChart(level, grid_type, element_type, hg_file,
                      wave_names, H_min, H_step, G_min, G_step,
                      output_dir):
    """
    保存hg文件的等值线分类数据，用于可视化同潮图
    
    :param level:               离散格网等级
    :param grid_type:           离散格网类型
    :param element_type:        离散格网元素类型
    :param hg_file:             调和分析得到的迟角振幅文件
    :param wave_names:          要获取的分潮H和/或G，如['M2_H', 'M2_G', 'S2_H', 'S2_G']
    :param H_min:               振幅的最小值
    :param H_step:              振幅的等值线步长
    :param G_min:               迟角的最小值
    :param G_step:              迟角的等值线步长
    :param output_dir:          同潮图数据输出目录
    """

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with open(hg_file, 'r') as hg_input:
        grid_hg = json.load(hg_input)

    prop_names = [f"{wave}_H" for wave in wave_names] + [f"{wave}_G" for wave in wave_names]
    grid_values = getHgData(level, grid_type, element_type, hg_file, prop_names)

    for prop_name in prop_names:
        if 'H' in prop_name: # for H
            grid_contour = Contour(level, grid_type, element_type,
                                   grid_values[prop_name], H_min, H_step)
        else:   # for G
            grid_contour = Contour(level, grid_type, element_type,
                                   grid_values[prop_name], G_min, G_step,
                                   is_cycle=True, cycle_bounds=(0, 360))

        if not grid_contour:
            print(f"无法生成'{prop_name}'的等值线，因为没有一个具有值的单元")
            return

        classified_grid_by_id, classified_grid_by_value = grid_contour.classifyGrid()
        contour_grid = grid_contour.calGridCountour((classified_grid_by_id, classified_grid_by_value))

        contour_features_class = {}
        for c_class, cell_ids in contour_grid.items():
            for cell_id in cell_ids:
                contour_features_class[cell_id] = c_class

        output_geojson = {}
        for key, value in grid_hg.items():
            if key != 'features':
                output_geojson[key] = value
            else:
                features = []
                for feature in grid_hg[key]:
                    feature_id = feature['properties']['codeId']
                    if feature_id not in contour_features_class:
                        continue

                    feature['properties'] = {
                        'class': contour_features_class[feature_id],
                        'codeId': feature_id
                    }
                    features.append(feature)
                output_geojson[key] = features
    
        with open(os.path.join(output_dir, f"level{level}_{prop_name}.geojson"), 'w', encoding='UTF-8') as output:
            json.dump(output_geojson, output, indent=2)

def main():

    saveCotideChart(level, grid_type, 
        element_type, hg_output_path, wave_names, 
        H_min, H_step, G_min, G_step, output_dir
    )


if __name__ =='__main__':
    print('start')
    main()
    print('end')
